/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2021-2021. All rights reserved.
 */

package org.apache.ranger.obs.server.http;

import com.sun.net.httpserver.Headers;
import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;

import org.codehaus.jackson.JsonFactory;
import org.codehaus.jackson.JsonGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.OutputStream;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Array;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import javax.management.AttributeNotFoundException;
import javax.management.InstanceNotFoundException;
import javax.management.IntrospectionException;
import javax.management.MBeanAttributeInfo;
import javax.management.MBeanException;
import javax.management.MBeanInfo;
import javax.management.MBeanServer;
import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;
import javax.management.ReflectionException;
import javax.management.RuntimeErrorException;
import javax.management.RuntimeMBeanException;
import javax.management.openmbean.CompositeData;
import javax.management.openmbean.CompositeType;
import javax.management.openmbean.TabularData;

public class JMXHttpHandler implements HttpHandler {

    private static final Logger LOG = LoggerFactory.getLogger(JMXHttpHandler.class);

    protected transient MBeanServer server = null;

    private ExecutorService threadPool = Executors.newFixedThreadPool(4);

    public JMXHttpHandler() {
        server = ManagementFactory.getPlatformMBeanServer();
    }

    @Override
    public void handle(final HttpExchange exchange) {
        this.threadPool.execute(new Runnable() {
            @Override
            public void run() {
                OutputStream outputStream = null;
                JsonGenerator generator = null;
                try {
                    Map<String, String> queries = new HashMap<String, String>();
                    String rawQuery = exchange.getRequestURI().getRawQuery();
                    if (null != rawQuery) {
                        String[] params = rawQuery.split("&");
                        for (String param : params) {
                            String[] paras = param.split("=");
                            if ("get".equals(paras[0])) {
                                queries.put("get", param.substring(4));
                            } else if ("qry".equals(paras[0])) {
                                queries.put("qry", param.substring(4));
                            }
                        }
                    }

                    Headers responseHeaders = exchange.getResponseHeaders();
                    responseHeaders.set("Content-Type", "application/json; charset=utf8");
                    exchange.sendResponseHeaders(200, 0);

                    outputStream = exchange.getResponseBody();
                    JsonFactory factory = new JsonFactory();
                    generator = factory.createJsonGenerator(outputStream);
                    generator.disable(JsonGenerator.Feature.AUTO_CLOSE_TARGET);
                    generator.useDefaultPrettyPrinter();
                    generator.writeStartObject();

                    if (server == null) {
                        generator.writeStringField("result", "ERROR");
                        generator.writeStringField("message", "MBeanServer could not be found");
                        generator.close();
                        LOG.error("MBeanServer could not be found.");
                        exchange.sendResponseHeaders(404, 0);
                    }

                    String method = queries.get("get");
                    if (method != null) {
                        String[] splits = method.split("\\:\\:");
                        if (splits.length != 2) {
                            generator.writeStringField("result", "ERROR");
                            generator.writeStringField("message", "query format is not as expected.");
                            generator.close();
                            exchange.sendResponseHeaders(400, 0);
                        }
                        listBeans(generator, new ObjectName(splits[0]), splits[1], exchange);
                        generator.close();
                    }

                    String query = queries.get("qry");
                    if (query == null) {
                        query = "*:*";
                    }
                    listBeans(generator, new ObjectName(query), null, exchange);
                } catch (IOException | MalformedObjectNameException e) {
                    LOG.error("processing jmx request", e);
                    try {
                        exchange.sendResponseHeaders(500, 0);
                    } catch (IOException e1) {
                        LOG.error("processing jmx request", e1);
                    }
                } finally {
                    try {
                        if (generator != null) {
                            generator.close();
                        }
                        if (outputStream != null) {
                            outputStream.close();
                        }
                    } catch (IOException e) {
                        LOG.error("processing jmx request", e);
                    }
                }
            }
        });
    }

    private void listBeans(JsonGenerator generator, ObjectName query, String attr, HttpExchange response)
        throws IOException {
        LOG.debug("query:" + query);
        Set<ObjectName> beanNames = null;
        beanNames = server.queryNames(query, null);

        generator.writeArrayFieldStart("beans");
        Iterator<ObjectName> iterator = beanNames.iterator();
        while (iterator.hasNext()) {
            ObjectName beanName = iterator.next();
            MBeanInfo bean;
            String className = "";
            Object attribute = null;
            try {
                bean = server.getMBeanInfo(beanName);
                className = bean.getClassName();
                String prs = "";
                try {
                    if ("org.apache.commons.modeler.BaseModelMBean".equals(className)) {
                        prs = "modelerType";
                        className = (String) server.getAttribute(beanName, prs);
                    }
                    if (attr != null) {
                        prs = attr;
                        attribute = server.getAttribute(beanName, prs);
                    }
                } catch (AttributeNotFoundException | MBeanException | RuntimeException | ReflectionException e) {
                    LOG.error("get attr " + prs + " of " + beanName, e);
                }
            } catch (InstanceNotFoundException | IntrospectionException | ReflectionException e) {
                LOG.error("get bean " + beanName, e);
                continue;
            }

            generator.writeStartObject();
            generator.writeStringField("name", beanName.toString());

            generator.writeStringField("modelerType", className);
            if (attr != null && attribute == null) {
                generator.writeStringField("result", "ERROR");
                generator.writeStringField("message", "attr with name " + attr + " was not found");
                generator.writeEndObject();
                generator.writeEndArray();
                generator.close();
                response.sendResponseHeaders(404, 0);
                return;
            }

            if (attr != null) {
                writeAttribute(generator, attr, attribute);
            } else {
                MBeanAttributeInfo[] attrs = bean.getAttributes();
                for (int i = 0; i < attrs.length; i++) {
                    writeAttribute(generator, beanName, attrs[i]);
                }
            }
            generator.writeEndObject();
        }
        generator.writeEndArray();
    }

    private void writeAttribute(JsonGenerator generator, ObjectName bean, MBeanAttributeInfo attrInfo)
        throws IOException {
        if (!attrInfo.isReadable()) {
            return;
        }
        String name = attrInfo.getName();
        if ("modelerType".equals(name)) {
            return;
        }
        if (name.indexOf("=") >= 0 || name.indexOf(":") >= 0 || name.indexOf(" ") >= 0) {
            return;
        }
        Object value = null;
        try {
            value = server.getAttribute(bean, name);
        } catch (RuntimeMBeanException e) {
            if (e.getCause() instanceof UnsupportedOperationException) {
                LOG.debug("get attr " + name + " of " + bean, e);
            } else {
                LOG.error("get attr " + name + " of " + bean, e);
            }
            return;
        } catch (AttributeNotFoundException | MBeanException | RuntimeException | ReflectionException
            | InstanceNotFoundException e) {
            LOG.error("get attr " + name + " of " + bean, e);
            return;
        }

        writeAttribute(generator, name, value);
    }

    private void writeAttribute(JsonGenerator generator, String name, Object value) throws IOException {
        generator.writeFieldName(name);
        writeObject(generator, value);
    }

    private void writeObject(JsonGenerator generator, Object value) throws IOException {
        if (value == null) {
            generator.writeNull();
        } else {
            Class<?> classObj = value.getClass();

            if (value instanceof Number) {
                Number num = (Number) value;
                generator.writeNumber(num.toString());
            } else if (value instanceof Boolean) {
                Boolean bool = (Boolean) value;
                generator.writeBoolean(bool);
            } else if (classObj.isArray()) {
                generator.writeStartArray();
                int len = Array.getLength(value);
                for (int j = 0; j < len; j++) {
                    Object item = Array.get(value, j);
                    writeObject(generator, item);
                }
                generator.writeEndArray();
            } else if (value instanceof CompositeData) {
                CompositeData data = (CompositeData) value;
                CompositeType type = data.getCompositeType();
                Set<String> keys = type.keySet();
                generator.writeStartObject();
                for (String k : keys) {
                    writeAttribute(generator, k, data.get(k));
                }
                generator.writeEndObject();
            } else if (value instanceof TabularData) {
                TabularData data = (TabularData) value;
                generator.writeStartArray();
                for (Object v : data.values()) {
                    writeObject(generator, v);
                }
                generator.writeEndArray();
            } else {
                generator.writeString(value.toString());
            }
        }
    }
}
